package cn.jinbyte.web.wrapper;


import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import lombok.Getter;
import org.jetbrains.annotations.NotNull;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

/**
 * 可重复读取响应内容包装器
 *
 * @author jinty
 */
public class RepeatReadResponseWrapper extends HttpServletResponseWrapper {
    private final ByteArrayOutputStream cachedOutputStream = new ByteArrayOutputStream();
    private final PrintWriter writer;
    @Getter
    private boolean isStreamingResponse;

    public RepeatReadResponseWrapper(HttpServletResponse response) {
        super(response);
        this.writer = new PrintWriter(this.cachedOutputStream, true);
        this.isStreamingResponse = false;
        this.checkStreamingResponse();
    }

    public void setContentType(String type) {
        super.setContentType(type);
        if (type != null) {
            String lowerType = type.toLowerCase();
            this.isStreamingResponse = lowerType.contains("text/event-stream");
        }

    }

    private void checkStreamingResponse() {
        String contentType = this.getContentType();
        if (contentType != null) {
            String lowerType = contentType.toLowerCase();
            this.isStreamingResponse = lowerType.contains("text/event-stream");
        }

    }

    public ServletOutputStream getOutputStream() throws IOException {
        this.checkStreamingResponse();
        return this.isStreamingResponse ? super.getOutputStream() : new ServletOutputStream() {
            public boolean isReady() {
                return true;
            }

            public void setWriteListener(WriteListener writeListener) {
            }

            public void write(int b) {
                RepeatReadResponseWrapper.this.cachedOutputStream.write(b);
            }

            public void write(byte @NotNull [] b) throws IOException {
                RepeatReadResponseWrapper.this.cachedOutputStream.write(b);
            }

            public void write(byte @NotNull [] b, int off, int len) {
                RepeatReadResponseWrapper.this.cachedOutputStream.write(b, off, len);
            }
        };
    }

    public PrintWriter getWriter() throws IOException {
        this.checkStreamingResponse();
        return this.isStreamingResponse ? super.getWriter() : this.writer;
    }

    public String getResponseContent() {
        if (!this.isStreamingResponse) {
            this.writer.flush();
            return this.cachedOutputStream.toString(StandardCharsets.UTF_8);
        } else {
            return null;
        }
    }

    public void copyBodyToResponse() throws IOException {
        if (!this.isStreamingResponse && this.cachedOutputStream.size() > 0) {
            this.getResponse().getOutputStream().write(this.cachedOutputStream.toByteArray());
        }

    }
}
